package Encrypt

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	"crypto/md5"
	"crypto/rand"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/sha512"
	"encoding/base64"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"lht-go-plugins/internal/core"
)

// 默认 Key（AES 128）
var defaultKey = []byte("lht_encrypt_key_")

func useKey(key []byte) []byte {
	if len(key) == 0 {
		return defaultKey
	}
	return key
}

type EncryptPlugin struct {
	*core.BasePlugin
}

func NewEncryptPlugin() *EncryptPlugin {
	return &EncryptPlugin{
		BasePlugin: core.NewBasePlugin("encrypt", 900),
	}
}

func (e *EncryptPlugin) Init(cfg any) error { return nil }

func (e *EncryptPlugin) Start() error {
	return e.BasePlugin.Start(e.handleData)
}

func (e *EncryptPlugin) handleData(data []byte) {
	fmt.Printf("[EncryptPlugin] recv data: %s\n", string(data))
}

// ========================================================
//
//	🔢 Base64
//
// ========================================================
func (e *EncryptPlugin) Base64Encode(data []byte) string {
	return base64.StdEncoding.EncodeToString(data)
}

func (e *EncryptPlugin) Base64Decode(s string) ([]byte, error) {
	return base64.StdEncoding.DecodeString(s)
}

// ========================================================
//
//	🧩 Hash（MD5 / SHA 家族）
//
// ========================================================
func (e *EncryptPlugin) MD5Hex(data []byte) string {
	sum := md5.Sum(data)
	return hex.EncodeToString(sum[:])
}

func (e *EncryptPlugin) SHA1Hex(data []byte) string {
	sum := sha1.Sum(data)
	return hex.EncodeToString(sum[:])
}

func (e *EncryptPlugin) SHA256Hex(data []byte) string {
	sum := sha256.Sum256(data)
	return hex.EncodeToString(sum[:])
}

func (e *EncryptPlugin) SHA512Hex(data []byte) string {
	sum := sha512.Sum512(data)
	return hex.EncodeToString(sum[:])
}

// ========================================================
//
//	🔐 HMAC（带 Key 的 Hash 签名）
//
// ========================================================
func (e *EncryptPlugin) HmacSHA256(data, key []byte) string {
	k := useKey(key)
	mac := hmac.New(sha256.New, k)
	mac.Write(data)
	return hex.EncodeToString(mac.Sum(nil))
}

func (e *EncryptPlugin) HmacMD5(data, key []byte) string {
	k := useKey(key)
	mac := hmac.New(md5.New, k)
	mac.Write(data)
	return hex.EncodeToString(mac.Sum(nil))
}

// ========================================================
//
//	🔒 AES-CBC（你已有 → 保留，加入默认 key）
//
// ========================================================
func (e *EncryptPlugin) AesEncrypt(data, key []byte) (string, error) {
	key = useKey(key)
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	data = pkcs7Padding(data, block.BlockSize())
	cipherText := make([]byte, len(data))

	iv := key[:block.BlockSize()]
	cipher.NewCBCEncrypter(block, iv).CryptBlocks(cipherText, data)

	return base64.StdEncoding.EncodeToString(cipherText), nil
}

func (e *EncryptPlugin) AesDecrypt(cipherTextBase64 string, key []byte) ([]byte, error) {
	cipherText, err := base64.StdEncoding.DecodeString(cipherTextBase64)
	if err != nil {
		return nil, err
	}

	key = useKey(key)
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	if len(cipherText)%block.BlockSize() != 0 {
		return nil, errors.New("cipherText is not a multiple of block size")
	}

	iv := key[:block.BlockSize()]
	plain := make([]byte, len(cipherText))
	cipher.NewCBCDecrypter(block, iv).CryptBlocks(plain, cipherText)

	return pkcs7UnPadding(plain), nil
}

// ========================================================
//
//	🚀 AES-GCM（推荐，安全且高效）
//
// ========================================================
func (e *EncryptPlugin) AesGCMEncrypt(data, key []byte) (string, error) {
	key = useKey(key)

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return "", err
	}

	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(rand.Reader, nonce); err != nil {
		return "", err
	}

	encrypted := gcm.Seal(nonce, nonce, data, nil)
	return base64.StdEncoding.EncodeToString(encrypted), nil
}

func (e *EncryptPlugin) AesGCMDecrypt(cipherTextBase64 string, key []byte) ([]byte, error) {
	cipherText, err := base64.StdEncoding.DecodeString(cipherTextBase64)
	if err != nil {
		return nil, err
	}

	key = useKey(key)

	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonceSize := gcm.NonceSize()
	nonce := cipherText[:nonceSize]
	data := cipherText[nonceSize:]

	return gcm.Open(nil, nonce, data, nil)
}

// ========================================================
//
//	⚠ AES-ECB（兼容旧系统）
//
// ========================================================
func (e *EncryptPlugin) AesECBEncrypt(data, key []byte) (string, error) {
	key = useKey(key)
	block, err := aes.NewCipher(key)
	if err != nil {
		return "", err
	}

	data = pkcs7Padding(data, block.BlockSize())
	encrypted := make([]byte, len(data))

	for i := 0; i < len(data); i += block.BlockSize() {
		block.Encrypt(encrypted[i:], data[i:])
	}

	return base64.StdEncoding.EncodeToString(encrypted), nil
}

func (e *EncryptPlugin) AesECBDecrypt(cipherTextBase64 string, key []byte) ([]byte, error) {
	cipherText, err := base64.StdEncoding.DecodeString(cipherTextBase64)
	if err != nil {
		return nil, err
	}

	key = useKey(key)
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	decrypted := make([]byte, len(cipherText))
	for i := 0; i < len(cipherText); i += block.BlockSize() {
		block.Decrypt(decrypted[i:], cipherText[i:])
	}

	return pkcs7UnPadding(decrypted), nil
}

// ========================================================
//
//	工具函数
//
// ========================================================
func pkcs7Padding(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	pad := bytesRepeat(byte(padding), padding)
	return append(data, pad...)
}

func pkcs7UnPadding(data []byte) []byte {
	if len(data) == 0 {
		return nil
	}
	unPadding := int(data[len(data)-1])
	return data[:len(data)-unPadding]
}

func bytesRepeat(b byte, count int) []byte {
	result := make([]byte, count)
	for i := range result {
		result[i] = b
	}
	return result
}
